package com.gy.ai.agent.platform.siliconflow;

import com.alibaba.fastjson2.JSON;
import com.gy.ai.agent.BasicPlatform;
import com.gy.ai.agent.client.BaseReqClient;
import com.gy.ai.agent.client.SseClient;
import com.gy.ai.agent.core.ReqContext;
import com.gy.ai.agent.core.ReqOptions;
import com.gy.ai.agent.enums.EncodingFormatEnums;
import com.gy.ai.agent.enums.RoleEnums;
import com.gy.ai.agent.listener.BaseReqClientListener;
import com.gy.ai.agent.listener.ReqClientListener;
import com.gy.ai.agent.listener.StreamRespListener;
import com.gy.ai.agent.platform.siliconflow.entity.*;
import com.gy.ai.agent.util.SysOkHttpClientUtil;
import lombok.extern.slf4j.Slf4j;

import java.math.BigDecimal;
import java.util.Collections;
import java.util.List;

/**
 * @author LGY
 * Create by 2025/6/16 11:36
 */
@Slf4j
public class SiliconflowPlatform extends BasicPlatform<SiliconflowPlatformConfig> {

    // 创建文本对话请求
    public static final String TEXT_CHART = "/chat/completions";
    // 创建文本向量请求 （嵌入请求）
    public static final String TEXT_EMBEDDINGS = "/embeddings";
    // 创建文本重排序请求
    public static final String TEXT_RERANK = "/rerank";
    // 创建文生图的请求
    public static final String TEXT_TO_IMAGE = "/images/generations";

    public SiliconflowPlatform(SiliconflowPlatformConfig config) {
        super(config);
    }

    @Override
    public SiliconflowChatResp chat(String prompt) {
        SiliconflowChatReq req = new SiliconflowChatReq();
        req.setMessages(Collections.singletonList(new SiliconflowChatReq.Message().setRole(RoleEnums.USER.getValue()).setContent(prompt)));
        return chat(req);
    }


    @Override
    public SiliconflowChatResp chat(ReqOptions options) {
        SiliconflowPlatformConfig siliconflowPlatformConfig = this.getConfig();
        String url = siliconflowPlatformConfig.getEndpoint() + TEXT_CHART;
        // 请求参数
        SiliconflowChatReq req = (SiliconflowChatReq) options;
        req.setModel(siliconflowPlatformConfig.getModel())
                .setStream(false);
        // 发起请求
        return SysOkHttpClientUtil.doPost(url, siliconflowPlatformConfig.getHeaders(), JSON.toJSONString(req), SiliconflowChatResp.class);
    }

    @Override
    public void chatStream(String prompt, StreamRespListener streamResponseListener) {
        SiliconflowChatReq req = new SiliconflowChatReq();
        req.setMessages(Collections.singletonList(new SiliconflowChatReq.Message().setRole(RoleEnums.USER.getValue()).setContent(prompt)));
        chatStream(req, streamResponseListener);
    }

    @Override
    public void chatStream(ReqOptions options, StreamRespListener streamResponseListener) {
        SiliconflowPlatformConfig siliconflowPlatformConfig = this.getConfig();
        BaseReqClient reqClient = new SseClient();
        String url = siliconflowPlatformConfig.getEndpoint() + TEXT_CHART;
        SiliconflowChatReq req = (SiliconflowChatReq) options;
        req.setModel(siliconflowPlatformConfig.getModel())
                .setStream(true);
        ReqClientListener listener = new BaseReqClientListener(this, streamResponseListener, reqClient, new ReqContext(this, reqClient));
        reqClient.start(url, siliconflowPlatformConfig.getHeaders(), JSON.toJSONString(req), listener);
    }

    @Override
    public SiliconflowEmbeddingResp embedding(String input) {
        SiliconflowEmbeddingReq req = new SiliconflowEmbeddingReq();
        req.setInput(input)
                .setEncoding_format(EncodingFormatEnums.FLOAT.getValue());
        return embedding(req);
    }

    @Override
    public SiliconflowEmbeddingResp embedding(ReqOptions options) {
        SiliconflowPlatformConfig siliconflowPlatformConfig = this.getConfig();
        String url = siliconflowPlatformConfig.getEndpoint() + TEXT_EMBEDDINGS;
        // 请求参数
        SiliconflowEmbeddingReq req = (SiliconflowEmbeddingReq) options;
        req.setModel(siliconflowPlatformConfig.getModel());
        // 发起请求
        return SysOkHttpClientUtil.doPost(url, siliconflowPlatformConfig.getHeaders(), JSON.toJSONString(req), SiliconflowEmbeddingResp.class);
    }

    @Override
    public SiliconflowRerankResp rerank(String query, List<String> documents) {
        SiliconflowRerankReq req = new SiliconflowRerankReq();
        req.setQuery(query)
                .setDocuments(documents)
                .setTop_n(5);
        return rerank(req);
    }

    @Override
    public SiliconflowRerankResp rerank(ReqOptions options) {
        SiliconflowPlatformConfig siliconflowPlatformConfig = this.getConfig();
        String url = siliconflowPlatformConfig.getEndpoint() + TEXT_RERANK;
        // 请求参数
        SiliconflowRerankReq req = (SiliconflowRerankReq) options;
        req.setModel(siliconflowPlatformConfig.getModel());
        // 发起请求
        return SysOkHttpClientUtil.doPost(url, siliconflowPlatformConfig.getHeaders(), JSON.toJSONString(req), SiliconflowRerankResp.class);
    }

    @Override
    public SiliconflowGenerateImageResp generateImage(String prompt) {
        SiliconflowGenerateImageReq req = new SiliconflowGenerateImageReq();
        req.setPrompt(prompt)
                .setBatch_size(1)
                .setImage_size("1024x1024")
                .setNum_inference_steps(20)
                .setGuidance_scale(BigDecimal.valueOf(7.5));
        return generateImage(req);
    }

    @Override
    public SiliconflowGenerateImageResp generateImage(ReqOptions options) {
        SiliconflowPlatformConfig siliconflowPlatformConfig = this.getConfig();
        String url = siliconflowPlatformConfig.getEndpoint() + TEXT_TO_IMAGE; // 请求参数
        SiliconflowGenerateImageReq req = (SiliconflowGenerateImageReq) options;
        req.setModel(siliconflowPlatformConfig.getModel());
        // 发起请求
        return SysOkHttpClientUtil.doPost(url, siliconflowPlatformConfig.getHeaders(), JSON.toJSONString(req), SiliconflowGenerateImageResp.class);
    }
}
